#!/usr/bin/env python
#--coding=utf-8--
import argparse
import logging
import os
import stat
import subprocess
import sys
import time
import requests

from socket import gethostname

"""
# 自动读取配置文件，并设置wifi

每分钟检查一次 是否有U盘挂载并且根目录有 
wifi.txt
第一行wifiSSID
第二行wifi 密码

读取信息 写入到 /etc/wpa_supplicant/wpa_supplicant.conf
并重启wlan

1. 安装依赖
usbmount 确保U盘自动挂载

sudo pip install requests
sudo apt-get install usbmount -y
Make sure it works in Stretch by changing MountFlags=slave to MountFlags=shared here:
sudo nano /lib/systemd/system/systemd-udevd.service

3. 添加定时任务
sudo crontab -e 
* *  * * *  /usr/local/bin/raspi_wificonf.py --api=http://api.host/wifi.conf --usb

TODO：
网络联通状态LED指示灯
"""

wpa_tpl ="""

ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1
country=GB

network={
    ssid="{SSID}"
    psk="{PASSWD}"
    key_mgmt=WPA-PSK
}

"""
class AutoWifiException(Exception):
    pass

def get_logger(logfilename, loggerName=None, colored=True):
    '''
    When a logger is created, the level is set to NOTSET (which causes all messages to be processed when the logger is the root logger, 
    or delegation to the parent when the logger is a non-root logger). 
    Note that the root logger is created with level WARNING.
    
    CRITICAL	50
    ERROR	40
    WARNING	30
    INFO	20
    DEBUG	10
    NOTSET	0
    
    '''
    stream_formatter = logging.Formatter('%(name)-4s: %(levelname)-4s %(message)s')    
    if colored:
        try:
            from colorlog import ColoredFormatter #window下必须 在 get_logger 之前导入 否则乱码
            LOGFORMAT = '%(log_color)s%(name)-4s%(reset)s: %(log_color)s%(levelname)-4s %(message)s%(reset)s'
            stream_formatter = ColoredFormatter(LOGFORMAT)
        except ImportError:
            pass
            
    rlog=logging.getLogger(loggerName)
    rlog.setLevel(logging.DEBUG)
    
    #stream
    stream = logging.StreamHandler(sys.stdout)
    stream.setFormatter(stream_formatter)
    stream.setLevel(logging.NOTSET)
    rlog.addHandler(stream)
    
    #file
    _file_handler = logging.FileHandler(logfilename,'a')
    _file_handler.setLevel(logging.NOTSET)
    _file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)-4s: %(levelname)s: %(message)s'))
    rlog.addHandler(_file_handler)
    return rlog

class Main(object):
    def __init__(self):
        self.wpa_filepath = "/etc/wpa_supplicant/wpa_supplicant.conf"
        self.conf_filepath = '/media/usb0/wifi.txt'
    
    def check_usb(self):
        if not os.path.isfile(self.conf_filepath):
            print('usbdrive not mount or %s not exist'%self.conf_filepath)
            return
        try:    
            conf = self.read_config()
        except AutoWifiException as e:
            rlog.error(e)
            return False
        suc = self.write_config(conf)
        if suc:
            self.reboot_wlan()

    def check_web(self, api):
        resp = requests.get(api, params={'host':gethostname()}, timeout = 3)
        if resp.status_code != 200:
            print('api: %s not valided, http code:%s'%(api, resp.status_code))
            return False
        mtime = 0    
        if resp.headers.get('Last-Modified'):
            _mtime = resp.headers.get('Last-Modified')
            mtime = time.mktime(time.strptime(_mtime, '%a, %d %b %Y %H:%M:%S %Z'))
            mtime += 3600 * 8 #上面返回的是GTM时区时间 这里暂时写死 增加8个小时
        _c = resp.text.split('\n')
        conf = {'ssid': _c[0].strip(), 'passwd':_c[1].strip()}
        if mtime:
            conf['mtime'] = mtime
        suc = self.write_config(conf)
        if suc:
            self.reboot_wlan()
            return True

    def read_config(self):
        """读取ssid 和密码
        """
        _st = os.stat(self.conf_filepath)
        st_mtime = _st[stat.ST_MTIME]
        with open(self.conf_filepath, 'r') as h:
            contents = h.read()
            try:
                _c = contents.split('\n')
                return {'ssid': _c[0].strip(), 'passwd':_c[1].strip(), 'mtime':st_mtime}
            except Exception:
                raise  AutoWifiException('please config file format')   
        return None

    def write_config(self, conf):
        """
        根据配置信息 写入wpa配置文件
        """
        _st = os.stat(self.wpa_filepath)
        st_mtime = _st[stat.ST_MTIME]
        if conf.get('mtime') and conf['mtime'] < st_mtime:
            print('wifi.txt not change')
            return False
        content = wpa_tpl.replace('{SSID}', conf['ssid']).replace('{PASSWD}', conf['passwd'])
        with open(self.wpa_filepath, 'w') as h:
            h.write(content)
            rlog.info('Write wifi success; ssid:%s'%conf['ssid'])
        return True

    def reboot_wlan(self):
        args = ['ifdown','wlan0']
        rlog.debug(subprocess.check_output(args, stderr=subprocess.STDOUT))
        args = ['ifup','wlan0']
        rlog.debug(subprocess.check_output(args, stderr=subprocess.STDOUT))
        rlog.info('wlan0 reboot success')

version='0.3'
parser = argparse.ArgumentParser(argument_default='-h')
parser.add_argument('-v','--version', action='version', version=version, help="show program's version number and exit")
parser.add_argument('--ssid', default="", nargs='?', help=u'wifi ssid')
parser.add_argument('--passwd', default="", nargs='?',help=u'wifi 密码')
parser.add_argument('--api',  default='', nargs='?', help=u'通过web api获取wifi信息')
parser.add_argument('--usb', default=False, action='store_true',help=u'检查U盘，获取配置信息')
parser.add_argument('--debug', default=False, action='store_true',help=u'window下调试')
args = parser.parse_args()

logfile = '/var/log/autowifi.log'
if args.debug:
    logfile = './autowifi.log'


rlog = get_logger(logfile)
main = Main()
#window下调试
if args.debug:
    main.conf_filepath = './wifi.txt'
    main.wpa_filepath = "./wpa_supplicant.conf"
if args.ssid:
    if not args.passwd:
        rlog.error('--passwd not set')
        sys.exit(1)
    main.write_config({'ssid':args.ssid, 'passwd':args.passwd})   
    if not args.debug:
        main.reboot_wlan()
    sys.exit(0)
    
if args.api:
    main.check_web(args.api)
if args.usb:    
    main.check_usb()

if not args.api and not args.usb:
    main.check_usb()
